import random
from collections import namedtuple, deque
import math

import itertools

Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward', 'done'))


class ReplayMemory(object):
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = deque([], capacity)

    def push(self, *args):
        transition = Transition(*args)
        self.memory.append(transition)

    def sample(self, batch_size):
        return random.sample(
            self.memory,
            min(batch_size, len(self.memory))
        )

    def __len__(self):
        return len(self.memory)


class PriorityReplayMemory(object):
    def __init__(self, capacity, alpha=0.5):
        self.alpha = alpha
        self.memory = deque([], math.floor(alpha * capacity))
        self.good_memory = deque([], math.floor(alpha * capacity))

    def push(self, *args):
        transition = Transition(*args)

        if transition.reward[0] == 1:
            self.good_memory.append(transition)
        else:
            self.memory.append(transition)

    def sample(self, batch_size):
        from_good = random.sample(self.good_memory, min(int(batch_size * self.alpha), len(self.good_memory)))
        from_bad = random.sample(self.memory, min(int(batch_size * (1 - self.alpha)), len(self.memory)))

        return from_good + from_bad

    def __len__(self):
        return len(self.memory) + len(self.good_memory)
